#!/usr/bin/python
# Author: Jimmy Saw
# Date: 02/12/2011
# This program extracts upstream sequences of a given COG identity from a
# completely sequenced genome.

import sys
#import re
from Bio import SeqIO

ptt_file = sys.argv[1]
seq_file = sys.argv[2]
cogid = sys.argv[3]

pf = open(ptt_file, "rU")
sf = open(seq_file, "rU")

lines = pf.readlines()

num = len(lines)

record = SeqIO.read(sf, "fasta")

#m1 = re.compile('^'cogid'')

i = 0

while i < num:
    if i == 3: #first record
        curr_line = lines[i].split('\t')
        curr_locustag = curr_line[5]
        curr_strand = curr_line[1]
        curr_coords = curr_line[0].split('..')
        curr_start = int(curr_coords[0])
        curr_stop = int(curr_coords[1])
        curr_cogid = curr_line[7]
        curr_def = curr_line[8]
        next_line = lines[i+1].split('\t')
        next_coords = next_line[0].split('..')
        next_start = int(next_coords[0])
        if curr_strand == "+":
            igs_start = 0
            igs_stop = curr_start - 1
            igs_seq = record.seq[igs_start:igs_stop]
        else:
            igs_start = curr_stop + 1
            igs_stop = next_start - 1
            igs_seq = record.seq[igs_start:igs_stop].reverse_complement()
        if curr_cogid == cogid:
            print ">" + curr_locustag #+ "\t" + curr_cogid + "\t" + str(igs_start) + ":" + str(igs_stop)
            print igs_seq
            #print curr_locustag, curr_cogid, igs_start, igs_stop, "IGS upstream =", igs_seq, "\n"

    elif 3 < i < num - 1:
        curr_line = lines[i].split('\t')
        curr_locustag = curr_line[5]
        curr_strand = curr_line[1]
        curr_coords = curr_line[0].split('..')
        curr_start = int(curr_coords[0])
        curr_stop = int(curr_coords[1])
        curr_cogid = curr_line[7]
        curr_def = curr_line[8]
        prev_line = lines[i-1].split('\t')
        prev_coords = prev_line[0].split('..')
        prev_start = int(prev_coords[0])
        prev_stop = int(prev_coords[1])
        next_line = lines[i+1].split('\t')
        next_coords = next_line[0].split('..')
        next_start = int(next_coords[0])
        next_stop = int(next_coords[1])
        if curr_strand == "+":
            igs_start = prev_stop + 1
            igs_stop = curr_start - 1
            igs_seq = record.seq[igs_start:igs_stop]
        else:
            igs_start = curr_stop + 1
            igs_stop = next_start - 1
            igs_seq = record.seq[igs_start:igs_stop].reverse_complement()

        if curr_cogid == cogid:
            print ">" + curr_locustag #+ "\t" + curr_cogid + "\t" + str(igs_start) + ":" + str(igs_stop)
            print igs_seq
            #print curr_locustag, curr_cogid, igs_start, igs_stop, "IGS upstream =", igs_seq, "\n"

    elif i == num - 1 : #last record
        curr_line = lines[i].split('\t')
        curr_locustag = curr_line[5]
        curr_strand = curr_line[1]
        curr_coords = curr_line[0].split('..')
        curr_start = int(curr_coords[0])
        curr_stop = int(curr_coords[1])
        curr_cogid = curr_line[7]
        curr_def = curr_line[8]
        prev_line = lines[i-1].split('\t')
        prev_coords = prev_line[0].split('..')
        prev_start = int(prev_coords[0])
        prev_stop = int(prev_coords[1])
        if curr_strand == "+":
            igs_start = prev_stop + 1
            igs_stop = curr_start - 1
            igs_seq = record.seq[igs_start:igs_stop]
        else:
            igs_start = curr_stop + 1
            igs_stop = len(record)
            igs_seq = record.seq[igs_start:igs_stop].reverse_complement()
        if curr_cogid == cogid:
            print ">" + curr_locustag #+ "\t" + curr_cogid + "\t" + str(igs_start) + ":" + str(igs_stop)
            print igs_seq
            #print curr_locustag, curr_cogid, igs_start, igs_stop, "IGS upstream =", igs_seq, "\n"
    i += 1
